# agent_config.py
import os

from langchain import hub
from langchain.agents import create_structured_chat_agent, AgentExecutor
from langchain.chains.combine_documents import create_stuff_documents_chain
from langchain.chains.retrieval import create_retrieval_chain
from langchain.memory import ConversationBufferMemory
from langchain_community.agent_toolkits import create_sql_agent
from langchain_community.utilities import OpenWeatherMapAPIWrapper
from langchain_core.prompts import ChatPromptTemplate, MessagesPlaceholder
from langchain_experimental.agents.agent_toolkits import create_python_agent
from langchain_experimental.tools import PythonREPLTool
from langchain_core.tools import tool, Tool

try:
    from ultra_ai.api.knowledge_rag import knowledge_retriever
except ImportError:
    from api.knowledge_rag import knowledge_retriever


def get_tools(model, selected_tools):
    tools = []

    @tool
    def placeholder(name: str) -> str:
        """这是一个占位工具，不要使用此工具。"""
        return "skip this step"

    @tool
    def file_retriever(question: str) -> str:
        """获取上传文件内容的工具, 当你需要文件内容来回答或总结时，请使用此工具，此工具会从文件中获取总结好的内容。

        "Use the following pieces of retrieved context to answer "
        "the question. If you don't know the answer,say that you "
        "don't know. Use three sentences maximum and keep the "
        "answer concise."
        "\n\n"
        {context}

        ####
        BEGIN!!!
        Question:
        """
        # 检查是否有文件被上传并且retriever已经初始化
        if 'retriever' in globals() and retriever is not None:

            # 检索
            docs = retriever.get_relevant_documents(question)

            # 提取相关片段内容
            retrieved_context = "\n\n".join([doc.page_content for doc in docs])

            # 如果检索到相关信息，构建并返回一个详细的回答提示
            if retrieved_context:
                prompt_content = """
                                "Use the following pieces of retrieved context to answer "
                                "the question. If you don't know the answer,say that you "
                                "don't know. Use three sentences maximum and keep the "
                                "answer concise."
                                "\n\n"
                                "{context}"
                        """
                prompt_template = ChatPromptTemplate.from_messages(
                    [
                        ("system", prompt_content),
                        ("human", "{input}")
                    ]
                )
                stuff_chain = create_stuff_documents_chain(llm=model, prompt=prompt_template)

                chain = create_retrieval_chain(retriever, stuff_chain)

                response = chain.invoke({"input": question})
                # return retrieved_context
                return response["answer"]
            else:
                return "未从上传的文件中找到相关信息，停止使用file_retrieve。"
        else:
            return "无文件被上传，停止使用file_retriever。"

    agent_python_executor = create_python_agent(
        llm=model,
        tool=PythonREPLTool(),
        verbose=True,
        agent_executor_kwargs={
            "handle_parsing_errors": True
        }
    )
    python_tool = Tool(
        name="python解释器",
        description="""当你需要借助Python解释器时，比如进行数学计算等操作时使用这个工具。
                        用自然语言把要求或者要求中的表达式提取给这个工具，生成要求对应的Python代码并返回代码执行的结果。""",
        func=agent_python_executor
    )

    os.environ["OPENWEATHERMAP_API_KEY"] = "879544e7845335ee2f4d3168fd81cd55"

    @tool
    def get_weather(city: str) -> str:
        """天气情况获取工具,当你需要获取天气时，请使用此工具，
        情况一：如果用户问题中没有给出具体的城市消息，那么请先让用户给出具体的城市再使用此工具
        Question：当前天气如何？
        Thought：用户没有给出查询哪个城市的信息
        Action：我需要知道您需要返回你想知道天气的城市名称！
        (Question/Thought/Action 重复调用，直到获取到最后的结果)、


        情况二：如果用户问题中的城市输入格式不符合使用条件，比如
        Question：广州天气条件如何
        Thought：用户现在查询广州的天气情况，但是输入为广州，而不是'GUANGZHOU,CN',所以不满足查询条件
        Action：将广州转变为GUANGZHOU,CN进行查询
        (Question/Thought/Action 重复调用，直到获取到最后的结果)、

        ####
        BEGIN!!!
        Question:
        """
        if "广州" == city:
            return "GUANGZHOU, CN"
        if "北京" == city:
            return "BEIJING, CN"
        if "上海" == city:
            return "SHANGHAI, CN"
        weather = OpenWeatherMapAPIWrapper()
        weather_data = weather.run(city)
        return weather_data

    # db = None
    # db_agent_executor = create_sql_agent(
    #     db=db, llm=model, agent_type="openai-tools", verbose=True, handle_parsing_errors=True
    # )
    # sql_tool = Tool(
    #     name="数据库查询工具",
    #     description="""
    #             当你本身无法获取到问题追定的数据信息的时候，你可以使用此数据库查询工具，将用户的问题转为SQL后，进行查询操作，最后工具会返回一个数据查询结果
    #         """,
    #     func=db_agent_executor.invoke
    # )

    if "Python解释器" in selected_tools:
        tools.append(python_tool)
    if "天气工具" in selected_tools:
        tools.append(get_weather)
    # if "数据库查询工具" in selected_tools:
    #     tools.append(sql_tool)

    tools.append(file_retriever)
    if selected_tools is None:
        tools.append(placeholder)
        # tools.append(placeholder)
        return tools

    return tools


def create_agent(model, tools_selection, uploaded_files, user_prompt):
    # 根据用户选择的工具初始化工具列表
    tools = get_tools(model, tools_selection)

    # 处理用户上传的知识库文件
    global retriever
    if uploaded_files:
        retriever = knowledge_retriever(uploaded_files)

    # 创建代理并加入记忆功能
    pulled_prompt = hub.pull("hwchase17/structured-chat-agent")
 #    react_prompt_content = """
 #        Respond to the human as helpfully and accurately as possible. You have access to the following tools:
 #
 #        {tools}
 #
 #        Use a json blob to specify a tool by providing an action key (tool name) and an action_input key (tool input).
 #        You are allow to use multiple tools in a single response.
 #        Your output should be related to {chat_history}.
 #
 #        Valid "action" values: "Final Answer" or {tool_names}
 #
 #        Provide only ONE action per $JSON_BLOB, as shown:
 #
 #        ```
 #        {{
 #          "action": $TOOL_NAME,
 #          "action_input": $INPUT
 #        }}
 #        ```
 #
 #        Follow this format:
 #
 #        Question: input question to answer
 #        Thought: consider previous and subsequent steps
 #        Action:
 #        ```
 #        $JSON_BLOB
 #        ```
 #        Observation: action result
 #        ... (repeat Thought/Action/Observation N times)
 #        Thought: I know what to respond
 #        Action:
 #        ```
 #        {{
 #          "action": "Final Answer",
 #          "action_input": "Final response to human"
 #        }}
 #
 #        Begin! Reminder to ALWAYS respond with a valid json blob of a single action. Use tools if necessary. Respond directly if appropriate. Format is Action:```$JSON_BLOB```then Observation"),
 #          ("placeholder", "{chat_history}"),
 #          ("human", "{input}
 #
 #        {agent_scratchpad}
 #         (reminder to respond in a JSON blob no matter what)")
 # """
 #
 #    react_prompt_content = user_prompt + react_prompt_content
 #
    # system = '''Respond to the human as helpfully and accurately as possible. You have access to the following tools:
    #
    # {tools}
    #
    # Use a json blob to specify a tool by providing an action key (tool name) and an action_input key (tool input).
    #
    # Valid "action" values: "Final Answer" or {tool_names}
    #
    # Provide only ONE action per $JSON_BLOB, as shown:
    #
    # ```
    # {{
    #   "action": $TOOL_NAME,
    #   "action_input": $INPUT
    # }}
    # ```
    #
    # Follow this format:
    #
    # Question: input question to answer
    # Thought: consider previous and subsequent steps
    # Action:
    # ```
    # $JSON_BLOB
    # ```
    # Observation: action result
    # ... (repeat Thought/Action/Observation N times)
    # Thought: I know what to respond
    # Action:
    # ```
    # {{
    #   "action": "Final Answer",
    #   "action_input": "Final response to human"
    # }}
    #
    # Begin! Reminder to ALWAYS respond with a valid json blob of a single action. Use tools if necessary. Respond directly if appropriate. Format is Action:```$JSON_BLOB```then Observation'''
    #
    # human = '''
    #
    # {input}
    #
    # {agent_scratchpad}
    #
    # (reminder to respond in a JSON blob no matter what)'''
    #
    # pulled_prompt = ChatPromptTemplate.from_messages(
    #     [
    #         ("system", system),
    #         MessagesPlaceholder(variable_name="chat_history"),
    #         ("human", human),
    #     ]
    # )



    agent = create_structured_chat_agent(
        llm=model,
        tools=tools,
        prompt=pulled_prompt
    )

    # memory = ConversationBufferMemory(
    #     return_history=True,
    #     memory_key="chat_history"
    # )
    memory = ConversationBufferMemory(
        return_messages=True,
        memory_key="chat_history"
    )

    agent_executor = AgentExecutor.from_agent_and_tools(
        tools=tools,
        agent=agent,
        memory=memory,
        verbose=True,
        handle_parsing_errors=True
    )

    return agent_executor
